#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/8/13 19:07
# @Author  : CHEN Wang
# @Site    :
# @File    : selection_backtest.py
# @Software: PyCharm

"""
脚本说明: 横截面选取资产策略的回测，支持股票，基金，行业，等； 不同于因子有效性检验，这里一般是对经过检验证明有效的因子做回测持续跟踪
"""


import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
from quant_researcher.quant.datasource_fetch.fund_api import api_func_old
from quant_researcher.quant.factors.factor_preprocess import preprocess
from quant_researcher.quant.datasource_fetch.fund_api import fund_nav_related
from quant_researcher.quant.datasource_fetch.stock_api import stock_price_related
from quant_researcher.quant.datasource_fetch.index_api import index_components, index_price_related
from quant_researcher.quant.performance_attribution.core_functions.performance_analysis.performance import simple_perf
from quant_researcher.quant.project_tool.common_var import PY_DATE_FMT_WITH_N_DASH
from quant_researcher.quant.project_tool.db_operator import db_conn
from quant_researcher.quant.project_tool import assert_tool
from quant_researcher.quant.project_tool.rebalance_date import get_rebalance_date

CON = db_conn.get_derivative_data_conn()
CON_factor = db_conn.get_tk_factors_conn()
CON_astk = db_conn.get_stock_conn()

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


class BackTest(object):  # todo 支持手续费设定，支持选取资产的组合方式（平均还是市值加权还是根据指标加权等）
    def __init__(self,
                 asset_type='stock',
                 sector_name: list = None,
                 sector_type='industry',
                 begin_date='',
                 end_date='',
                 rank_method=preprocess.cs_percentile_rank,
                 weight_method='等权',
                 benchmark=None,
                 indi_data=None,
                 indi_direction=None,
                 indi_screen=None,
                 abnormal_threshold=0.11,
                 freq='M',
                 period=1,
                 commission=0.003):
        """
        资产筛选回测类
        :param asset_type: 资产类型，支持股票，基金，行业和其他类型资产（有些分散在不同表，还需实现）
        :param sector_name: 板块名称，目前仅适用于股票，list类型，如特定行业、概念或指数代码
        :param sector_type: 板块种类，目前仅适用于股票和行业，当资产类型为股票时，支持行业，概念和指数；
                            当资产类型为行业时，支持申万，中证等行业
        :param begin_date: 回测开始日期，留空则默认为指标数据开始日期
        :param end_date: 回测结束日期，留空则默认为指标数据结束日期
        :param rank_method: 指标排序方法，简单排序或其他方法
        :param weight_method: 所选资产权重方法，支持等权，市值加权
        :param benchmark: 回测基准, 可以输入指数代码，如'000001'；输入avg代表资产池平均，
        :param indi_data: 指标数据，需包含code，tradedate，和指标列（支持多个指标），
        一个包含ROE指标的数据形状如下：
        —————————————
        code    |tradedate   |roe
        000001|2010-01-03| 0.2
        000002|2010-01-03| 0.15
        :param indi_direction: 指标方向，正向为1，负向为0，如一正一负两个指标为[1,0]。
        :param indi_screen: 指标筛选数据，和indi_data形状相同，符合条件的数据数值为1，不符合条件的数值为0
        ———————————————
        code    |tradedate   |in_out
        000001|2010-01-01| 1           # 符合条件的数据
        000002|2010-01-01| 0           # 不符合条件的数据
        :param abnormal_threshold: 日收益数据异常阈值
        :param freq: 再平衡频率（调仓频率）支持pandas提供的频率，‘M’频率为月末调仓，'W-FRI'为周五调仓，其他频率
        详情参见 https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases
        :param int, period: 再平衡周期，是每几个月，或者每几周
        """
        assert 'tradedate' in indi_data.columns, "请确认数据包含tradedate列"
        assert 'code' in indi_data.columns, "请确认数据包含code列"
        self.asset_type = asset_type
        self.weight_method = weight_method
        self.indi_data = indi_data[(indi_data.tradedate >= begin_date) & (indi_data.tradedate <=end_date)]
        self.begin_date, self.end_date = self.get_begin_end(begin_date, end_date)
        if asset_type == 'stock':
            type_list = ['industry', 'concept', 'index']
            assert_tool.value_in_list(sector_type, type_list, name='板块类型')
            assert_tool.value_is_instance(sector_name, list)
            # assert isinstance(sector_name, list), "请确认sector_name类型为list，单一板块为请输入['板块名称'],指数暂不支持混合股票池"
            self.sector_name = sector_name
            self.mix_pool = len(sector_name) > 1
            self.sector_type = sector_type
            print(f"本次分析的是{self.sector_name}")
        if asset_type == 'industry':
            type_list = ['申万一级', '中证800一级', '中证全指一级']
            assert_tool.value_in_list(sector_type, type_list, name='板块类型')
        self.benchmark = benchmark
        self.bench_code = self.get_bench_pool()
        self.rank_method = rank_method
        self.indi_direction = indi_direction
        self.indi_screen = indi_screen
        self.asset_pool = self.get_asset_pool()
        if weight_method == '市值加权':
            assert self.asset_type=='stock', "只有股票资产支持市值加权"
            self.stock_mkv = api_func_old.get_stock_mkv(self.begin_date, self.end_date, self.asset_pool, CON_astk)
        self.abnormal_threshold = abnormal_threshold
        self.freq = freq
        self.period = period
        self.rebalance_freq = str(self.period) + str(self.freq)
        self.rebalance_date = self.get_rebalance_date()
        self.commission = commission
        self.indicators_lst = self.indi_data.columns.difference(['tradedate', 'code']).tolist()
        assert len(self.indi_direction) == len(self.indicators_lst), \
            f"请确认标记所有指标的方向（指标个数:{len(self.indicators_lst)}, 已标记方向个数:{len(self.indi_direction)})"
        # self.stock_label = self.get_stock_label()
        print("获取资产池行情")
        self.asset_ret = self.get_asset_ret()
        print("获取基准行情")
        self.bench_ret = self.get_benchmark()
        print("更新有效资产池")
        self.update_stock_pool()

    def get_begin_end(self, begin, end):
        indicator_min_date = self.indi_data['tradedate'].min()
        indicator_max_date = self.indi_data['tradedate'].max()
        valid_begin = max(pd.to_datetime(begin), indicator_min_date).strftime(PY_DATE_FMT_WITH_N_DASH)
        valid_end = min(pd.to_datetime(end), indicator_max_date).strftime(PY_DATE_FMT_WITH_N_DASH)
        return valid_begin, valid_end

    def get_rebalance_date(self):
        date = get_rebalance_date(self.begin_date, self.end_date,
                                       freq=self.freq, period=self.period,
                                        business_day=True,
                       initial_start=True,
                       initial_end=True)

        rebalance_date = list(pd.to_datetime(date))
        return rebalance_date

    def update_data(self, data, direction):
        self.indi_data = data
        self.rebalance_date = self.get_rebalance_date()
        self.indicators_lst = self.indi_data.columns.difference(['tradedate', 'code']).tolist()
        self.indi_direction = direction

    def get_candidates(self, top=3, pct=False):  # todo 支持多种筛选方式
        mask = - pd.Series(self.indi_direction).astype(bool)
        if self.indi_screen:
            screen_data = self.indi_screen.set_index(['code', 'tradedate']).iloc[:, -1]
            data = self.indi_data.set_index(['code', 'tradedate'])[screen_data > 0]
        else:
            data = self.indi_data.set_index(['code', 'tradedate']).dropna()
        data = data[data.index.get_level_values(0).isin(self.asset_pool)]
        ranks = self.rank_method(data)
        ranks[ranks.columns[mask]] = 1 - ranks[ranks.columns[mask]]
        ranks = ranks.mean(axis=1).groupby('tradedate').rank(pct=True)
        ranks = ranks[~ranks.index.duplicated()]
        date_range = pd.date_range(self.begin_date, self.end_date)
        # 选取rebalance_date与date_range的交集
        rebalance_date = set(self.rebalance_date).intersection(set(date_range.tolist()))
        ranks = ranks.unstack(level=0).reindex(date_range).loc[rebalance_date].T.stack()  # 未进行前值填充，避免跨期比较
        if pct:
            candidates = ranks[(ranks > 1 - top / 100) & (ranks <= 1)].reset_index()
        else:
            ranks = ranks.groupby(level=1).rank(ascending=False)
            candidates = ranks[ranks <= top].reset_index()
        candidates.columns = ['code', 'tradedate', 'score']
        return candidates

    def get_asset_pool(self):
        external_pool = self.indi_data['code'].unique().tolist()
        if self.asset_type == 'stock':
            if self.sector_type == 'industry':
                stock_pool_df = index_components.get_index_hist_stock_pool(self.sector_name, self.begin_date, self.end_date)
                stock_pool_df = stock_pool_df.rename(columns={'stock_code': 'stockcode'})
            elif 'concept' == self.sector_type:
                stock_pool_df = index_components.get_concept_hist_stock_pool(self.sector_name)
            elif 'index' == self.sector_type:
                stock_pool_df = index_components.get_index_hist_stock_pool(self.sector_name, self.begin_date, self.end_date)
                stock_pool_df = stock_pool_df.rename(columns={'stock_code': 'stockcode'})
            temp = stock_pool_df['stockcode'].tolist()
        # elif self.asset_type == 'fund':
        #     temp = data_dealer.get_fund_pool(self.sector_name)
        # elif self.asset_type == 'industry':
        #     temp = data_dealer.get_industry_pool(self.sector_name)
        else:
            temp = external_pool
        if len(external_pool) < len(temp):
            print('指标数据资产池小于对应类别资产池')
        return list(set(external_pool).intersection(set(temp)))

    # def get_stock_label(self):
    #     result = pd.read_sql("select distinct Name,STOCKCODE "
    #                          "from longbo.jl_base_1102 "
    #                          f"where STOCKCODE in {tuple(self.stock_pool)} "
    #                          f"and F1102_032 in ('90','83')"
    #                          f"and F1102_022 = 1", CON_wst)
    #     return result

    def get_bench_pool(self):
        if self.benchmark == 'avg':
            bench_code = ['avg']
        else:
            bench_code = ['000001'] if self.benchmark is None else [self.benchmark]
        if self.asset_type == 'stock':
            if 'industry' == self.sector_type:
                pool = pd.read_sql("select index_code, substr(index_sname,5,length(index_sname)-4) as cname "
                                   "from idx_baseinfo "
                                   "where left(index_sname,4) = '申万一级'", CON)
                assert_tool.all_value_in_list(self.sector_name, pool.cname.tolist(), name='行业名称', throw_it=False)
                bench_code = pool[pool.cname.isin(self.sector_name)]['index_code'].tolist()
            elif 'concept' == self.sector_type:
                pool = pd.read_sql('select distinct sector_name '
                                   'from astk_qt_sectorinfo  '
                                   "where sector_type = 2 ", CON_astk)
                assert_tool.all_value_in_list(self.sector_name, pool.cname.tolist(), name='概念名称', throw_it=False)
                bench_code = ['000001'] if self.benchmark is None else [self.benchmark]
            elif 'index' == self.sector_type:
                pool = pd.read_sql('select distinct index_code from idx_components_weight ', CON)
                if self.sector_name == ['AllA']:
                    bench_code = ['000001']
                else:
                    assert_tool.all_value_in_list(self.sector_name, pool.index_code.tolist(), name='指数名称', throw_it=False)
                    bench_code = self.sector_name
        return bench_code

    def update_stock_pool(self):
        self.asset_pool = list(set(self.asset_pool) & set(self.asset_ret.columns))
        # lst = ['600705','000863']
        # for abnormal in lst:
        #     if abnormal in self.stock_pool:
        #         self.stock_pool.remove(abnormal)

    def get_asset_ret(self):
        if self.asset_type == 'stock':
            temp = stock_price_related.get_stock_quote(self.asset_pool, self.begin_date, self.end_date,
                                                       select=['stock_code', 'end_date',
                                                                  'high', 'low', 'close as fap_close'])
            temp = temp.rename(columns={'stock_code': 'stockcode', 'end_date': 'tradedate'})
        elif self.asset_type == 'fund':
            # temp = fund_nav_related.get_fund_adj_nav(self.asset_pool, [('nav_date', self.begin_date, self.end_date)],
            #                                          return_only_nav=False)
            temp = fund_nav_related.get_fund_nav(fund_code=self.asset_pool, nav_type='adj', start_date=self.begin_date, end_date=self.end_date)
            temp['end_date'] = pd.to_datetime(temp['end_date'])
            temp['end_date'] = temp['end_date'].dt.date
            temp = temp.rename(columns={'end_date': 'tradedate', 'fund_code': 'fundcode', 'adjusted_nav': 'fap_close'})
        elif self.asset_type == 'industry':
            # temp1 = index_price_related.get_stock_index_quote(self.asset_pool, self.begin_date, self.end_date,
            #                                                  select=['trade_date as tradedate',
            #                                                          'index_code as industrycode',
            #                                                          'close as fap_close', 'high', 'low'])
            temp = index_price_related.get_stock_index_quote(
                self.asset_pool, self.begin_date, self.end_date,
                select=['trade_date as tradedate', 'index_code as industrycode', 'close as fap_close', 'high', 'low'])
        temp = temp.rename(columns={'stockcode': 'code', 'fundcode': 'code'})
        temp['tradedate'] = pd.to_datetime(temp['tradedate'])
        ret = temp.set_index(['tradedate', 'code'])['fap_close'].unstack().ffill().pct_change()
        ret = ret.where(ret.abs() < self.abnormal_threshold, 0)
        return ret

    def get_stock_mkv(self):
        mkv_df = pd.read_sql("select trade_date as tradedate, sec_code as code,"
                             "freefloat_mktval "
                             "from astk_fn_daily_valuation "
                             f"where trade_date >= '{self.begin_date.replace('-', '')}'"
                             f"and trade_date <= '{self.end_date.replace('-', '')}'"
                             f"and sec_code in {tuple(self.asset_pool)}", CON_astk)
        mkv_df = mkv_df.set_index(['tradedate', 'code'])['freefloat_mktval'].unstack()
        mkv_df.index = pd.to_datetime(mkv_df.index)
        return mkv_df

    def get_benchmark(self):
        if self.bench_code == ['avg']:
            bench_ret = self.asset_ret.mean(axis=1)
        else:
            benchmark = pd.read_sql("select trade_date as tradedate, index_code, close "
                                    "from idx_mkt_quote "
                                    f"where trade_date >= '{self.begin_date.replace('-', '')}' "
                                    f"and trade_date <= '{self.end_date.replace('-', '')}'"
                                    f"and index_code in {tuple(self.bench_code + [''])} ", CON)
            benchmark['tradedate'] = pd.to_datetime(benchmark['tradedate'])
            bench_ret = benchmark.set_index(['tradedate', 'index_code'])['close'].unstack().ffill().pct_change()
            if self.asset_type == 'stock':
                if self.mix_pool:
                    bench_ret = bench_ret.mean(axis=1)
                else:
                    bench_ret = bench_ret.iloc[:, -1]
            else:
                bench_ret = bench_ret.iloc[:, -1]
            bench_ret = bench_ret.where(bench_ret.abs() < 0.11, 0)
        return bench_ret

    def get_equity(self, candidates, stock_ret):
        portfolio_list = []
        for idx, date in enumerate(self.rebalance_date[:-1]):
            print(f'调仓日期:{date}')
            stocks = candidates.code[candidates['tradedate'] == date]
            period_ret = stock_ret[stocks].loc[date:self.rebalance_date[idx + 1]]
            if self.weight_method=='市值加权':
                weight_df = self.stock_mkv[stocks].loc[date:self.rebalance_date[idx + 1]]
                period_ret = (period_ret * weight_df / weight_df.sum(axis=1)[0]).sum(axis=1)
            else:
                period_ret = period_ret.mean(axis=1)
            period_ret.iloc[0] = 0  # 策略无法取得调仓日当天收益
            portfolio_list.append(period_ret)
        portfolio_ret = pd.concat(portfolio_list).replace([np.inf, -np.inf], 0)
        portfolio_ret = portfolio_ret[~portfolio_ret.index.duplicated()]
        portfolio_ret = portfolio_ret.where(abs(portfolio_ret) < self.abnormal_threshold, 0)
        return portfolio_ret

    def plot(self, result):
        result = result.ffill()
        active = result.eval('Portfolio-Benchmark')
        fig, ax = plt.subplots(2, 1, figsize=(8, 9))
        # if len(self.sector_name) == 1:
        #     title = f'{self.sector_name}板块精选股票业绩表现'
        # else:
        #     title = f'大消费板块精选股票业绩表现'
        title = '策略净值曲线'
        result.plot(ax=ax[0], grid=True, title=title, sharex=True, fontsize=14, figsize=(16, 9))
        # result.plot(ax=ax[0], grid=True, title=f'{self.sector_name}板块精选股票业绩表现{self.indicators_lst[0]}', sharex=True)
        active.plot(ax=ax[1], grid=True, title=f'超额收益', sharex=True, fontsize=14, figsize=(16, 9))
        ax[0].yaxis.set_major_formatter(mtick.PercentFormatter(1))
        ax[1].yaxis.set_major_formatter(mtick.PercentFormatter(1))
        # plt.savefig(f'D:/概念精选/{self.sector_name}.png')
        plt.show()

    def run(self):
        print("获取调仓记录")
        candidates = self.get_candidates()
        print("开始回测")
        portfolio_ret = self.get_equity(candidates, self.asset_ret)
        result = pd.concat([portfolio_ret.rename('Portfolio'), self.bench_ret.rename('Benchmark')], axis=1, sort=True)
        result = result.loc[result['Portfolio'].first_valid_index():]
        result = (result + 1).cumprod().ffill()
        commission_mutiplier_dict = {date:(1 - self.commission) ** (i+1) for i, date in enumerate(self.rebalance_date)}
        commission_mutiplier = pd.Series(commission_mutiplier_dict)
        commission_mutiplier = commission_mutiplier.reindex(commission_mutiplier.index.union(result.index)).ffill().reindex(result.index)
        result['Portfolio'] = result['Portfolio'] * commission_mutiplier
        result = result.div(result.bfill().iloc[0])  # 净值起点归1
        result.index = pd.to_datetime(result.index)
        # 不同资产类型交易日不同，单独计算基准的业绩表现
        portfolio_perf = simple_perf(portfolio_ret).rename('Portfolio')
        bench_perf = simple_perf(self.bench_ret).rename('Benchmark')

        annual_active = portfolio_perf.loc['年化收益率'] - bench_perf.loc['年化收益率']
        perf = pd.concat([portfolio_perf, bench_perf], axis=1).T
        print("开始作图")
        self.plot(result)
        print("回测完成")
        print(perf)
        print(f'年化超额收益率：{annual_active * 100:.4f}%')
        print("------------------------------------------")
        return result, candidates

    def transactions(self, result, candidates):
        candidates['portfolio_count'] = candidates.groupby('tradedate')['score'].transform('count')
        result1 = candidates.merge(result, on='tradedate', how='left')
        result1 = result1.sort_values(by='tradedate')
        result1['Portfolio'] = result1['Portfolio'] * 1e6
        result1 = result1.rename(columns={'Portfolio': 'value'})
        result1['value'] = result1['value'] / result1['portfolio_count']
        buy = result1.loc[:, ['code', 'tradedate', 'value']]
        buy['direction'] = '买进'
        sell = buy.replace('买进', '卖出')
        sell['code'] = sell['code'].shift()
        sell = sell.dropna()
        result = pd.concat([buy, sell]).sort_values(by='tradedate')
        result.to_excel(f"D:\{self.asset_type}_{self.begin_date}_{self.end_date}_{'+'.join(self.indicators_lst)}.xlsx")


if __name__ == '__main__':
    pass
    # pool = indis.sector_pool['cname']
    # active_dict = dict()
    # candidates_dict = dict()
    # for sector in pool[:]:
    #     indis = BackTest(sector_name=sector,
    #                      sector_type='industry',
    #                      begin_date=period_start,
    #                      end_date=period_end,
    #                      indi_data=indicators, indi_direction=[1])
    #     # indis.update_data(indicators)
    #     result, candidates = indis.run()
    #     active = result.eval('Portfolio-Benchmark').iloc[-1]
    #     active_dict[sector] = active
    #     candidates_dict[sector] = candidates
    # %%
    # mkv_df = pd.read_sql("select trade_date as tradedate, sec_code as stockcode,"
    #                      "total_mktval, freefloat_mktval "
    #                      "from astk_fn_daily_valuation "
    #                      f"where trade_date = '20200217' "
    #                      f"order by sec_code, trade_date", CON_a_stk)
    # bigconsume = mkv_df[mkv_df.stockcode.isin(indis.stock_pool)]
    # mkv_df = mkv_df.set_index(['tradedate','stockcode'])
    # mkv_df = mkv_df[~mkv_df.index.duplicated()]
    # test = mkv_df['total_mktval'].rename('全A')
    # bigconsume = bigconsume.set_index(['tradedate','stockcode'])
    # bigconsume = bigconsume[~bigconsume.index.duplicated()]
    # test2 = bigconsume['total_mktval'].rename('大消费')
# %%
# indicators = ['roe','pe','gross_profit_margin','sales_growth','cfo_to_earning','interest_cover']
# direction = [1,0,1,1,1,1]
# a_list = []
# for indi,direc in zip(indicators,direction):
#     indicators = pd.read_sql("select stockcode, reportdate,  "
#                              f"{indi} "
#                              "from tk_mfdb_23.stk_comp_key_financials "
#                              f"where reportdate >= '{period_start}'"
#                              f"and reportdate <= '{period_end}'"
#                              "order by reportdate, stockcode", CON)
#     indicators['reportdate'] = pd.to_datetime(indicators['reportdate'])
#     indicators = indicators.drop_duplicates(['stockcode', 'reportdate'])
#     indis.update_data(indicators,[direc])
#     a, b = indis.run()
#     a_list.append(a.eval("Portfolio-Benchmark").rename(f'{indi}'))
# output = pd.concat(a_list, axis=1)
# ax=output.plot(figsize=(12,7),grid=True,title='单一指标筛选超额收益对比',fontsize=14)
# ax.yaxis.set_major_formatter(mtick.PercentFormatter(1))
# plt.show()
# benchmark = pd.read_sql("select trade_date as tradedate, index_code, close "
#                         "from idx_mkt_quote "
#                         f"where index_code ='000001' "
#                         f"order by tradedate", CON)
